#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *
from datetime import datetime, UTC

parser = argparse.ArgumentParser(description='Plot nonhyrostatic reduced sphere mountain wave test results')
parser.add_argument('-i', '--input', help='Input data file', default='sc.720x360x40.h0.nc')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-o', '--output', help='Output figure')
parser.add_argument('-n', '--no-update', help='Do not update', action='store_true')
parser.add_argument('-s', '--update-seconds', help='Update interval in seconds', default=10)
args = parser.parse_args()

fig = plt.figure(figsize=(12, 6))
fig.canvas.mpl_connect('close_event', window_closed)
last_update_time = None
min_lon = -40
max_lon = 40
min_lat = -40
max_lat = 40
zlev = np.array([5000]) * units.m

def update():
	global last_update_time
	modified_time = os.path.getmtime(args.input)
	if last_update_time != None:
		if last_update_time < modified_time:
			last_update_time = modified_time
			need_update = True
		else:
			need_update = False
	else:
		last_update_time = modified_time
		need_update = True

	if not need_update: return

	fig.clf()

	f = xr.open_dataset(args.input)
	time = datetime.fromtimestamp(f.time.isel(time=args.time_step).values.tolist() / 1e9, UTC)

	fig.suptitle(f'Time: {time.strftime("%H:%M")}')

	z = f.gz_lev[args.time_step,:,:] * units.m**2 / units.s**2 / g
	w_5km = vinterp_z(z, f.w_lev[args.time_step,:,:], zlev)

	z = f.gz[args.time_step,:,:] * units.m**2 / units.s**2 / g
	qr_5km = vinterp_z(z, f.q_qr[args.time_step,:,:], zlev) * 1000

	levels = np.linspace(-10, 26, 19)
	colors = ListedColormap([
	  [0.58, 0.75, 0.85],
	  [0.73, 0.86, 0.89],
	  [0.73, 0.87, 0.91],
	  [0.89, 0.96, 1.00],
	  [1.00, 1.00, 1.00],
	  [1.00, 1.00, 1.00],
	  [0.98, 0.89, 0.89],
	  [0.95, 0.77, 0.71],
	  [0.95, 0.77, 0.70],
	  [0.92, 0.65, 0.51],
	  [0.93, 0.64, 0.49],
	  [0.90, 0.51, 0.35],
	  [0.86, 0.36, 0.20],
	  [0.87, 0.36, 0.22],
	  [0.75, 0.23, 0.15],
	  [0.62, 0.18, 0.09],
	  [0.62, 0.18, 0.09],
	  [0.46, 0.13, 0.06],
	])

	ax = fig.add_subplot(1, 2, 1, projection=ccrs.PlateCarree(central_longitude=0))
	ax.set_extent([min_lon, max_lon, min_lat, max_lat], crs=ccrs.PlateCarree())
	plot_contour_map(ax, w_5km[0,:,:], cmap=colors, levels=levels, contour_levels=levels[np.abs(levels)>=1])
	ax.set_title(f'w (m/s) @ 5km, {w_5km.min().values:.1f}, {w_5km.max().values:.1f}')

	levels = np.linspace(1, 15, 15)
	colors = ListedColormap([
		[1.00, 1.00, 1.00],
		[1.00, 0.96, 0.80],
		[1.00, 0.96, 0.80],
		[0.99, 0.90, 0.44],
		[1.00, 0.80, 0.20],
		[1.00, 0.80, 0.20],
		[1.00, 0.68, 0.20],
		[1.00, 0.60, 0.20],
		[1.00, 0.43, 0.20],
		[1.00, 0.43, 0.20],
		[0.98, 0.33, 0.00],
		[0.88, 0.16, 0.10],
		[0.89, 0.15, 0.11],
		[0.77, 0.12, 0.07],
	])

	ax = fig.add_subplot(1, 2, 2, projection=ccrs.PlateCarree(central_longitude=0))
	ax.set_extent([min_lon, max_lon, min_lat, max_lat], crs=ccrs.PlateCarree())
	plot_contour_map(ax, qr_5km[0,:,:], cmap=colors, levels=levels)
	ax.set_title(f'qr (g/kg) @ 5km {qr_5km.min().values:.1f}, {qr_5km.max().values:.1f}')

	f.close()

	plt.tight_layout()
	plt.draw()

update()

if args.output:
	plt.savefig(args.output)
	exit(0)
elif not args.no_update:
	while True:
		plt.pause(args.update_seconds)
		try:
			update()
		except:
			pass
else:
	plt.show()
